# 遗传算法求解旅行商问题（TSP）优化版
# 主要改进：坐标数据结构、交叉算子优化、精英保留策略、可视化增强

import numpy as np
import matplotlib.pyplot as plt
import random
from math import radians, sin, cos, sqrt, atan2

# 城市坐标数据重构（使用numpy数组提升性能）
city_names = ["北京", "上海", "广州", "深圳", "成都"]
cities = np.array([
    (116.4074, 39.9042),  # 北京
    (121.4737, 31.2304),  # 上海
    (113.2644, 23.1291),  # 广州
    (114.0579, 22.5431),  # 深圳
    (104.0668, 30.6706)   # 成都
])

# 遗传算法参数优化
POP_SIZE = 100    # 种群规模
GEN_MAX = 200     # 最大迭代次数
CROSS_RATE = 0.9  # 提高交叉概率
MUT_RATE = 0.2    # 适当提高变异概率
ELITE_RATE = 0.1  # 精英保留比例


# 初始化种群
def init_population(pop_size, city_num):
    population = []
    for _ in range(pop_size):
        individual = list(range(city_num))
        random.shuffle(individual)
        population.append(individual)
    return population


# 计算路径总距离（适应度函数）
def calculate_distance(individual):
    total_distance = 0
    for i in range(len(individual)):
        city_a = individual[i]
        city_b = individual[(i + 1) % len(individual)]
        total_distance += np.linalg.norm(cities[city_a] - cities[city_b])
    return 1 / total_distance  # 适应度取倒数（距离越短适应度越高）


# 选择操作（轮盘赌）
def selection(population, fitnesses):
    total_fitness = sum(fitnesses)
    pick = random.uniform(0, total_fitness)
    current = 0
    for i, fitness in enumerate(fitnesses):
        current += fitness
        if current > pick:
            return population[i]


# 交叉操作（顺序交叉）
def crossover(parent1, parent2):
    child = [-1] * len(parent1)
    start, end = sorted(random.sample(range(len(parent1)), 2))
    child[start:end] = parent1[start:end]
    for i in range(len(parent2)):
        if parent2[i] not in child:
            for j in range(len(child)):
                if child[j] == -1:
                    child[j] = parent2[i]
                    break
    return child


# 变异操作（交换变异）
def mutate(individual):
    if random.random() < MUT_RATE:
        idx1, idx2 = random.sample(range(len(individual)), 2)
        individual[idx1], individual[idx2] = individual[idx2], individual[idx1]


# 可视化函数
def visualize(generation, best_individual, best_fitness):
    plt.figure(figsize=(10, 5))

    # 绘制城市位置
    plt.subplot(1, 2, 1)
    plt.scatter(cities[:, 0], cities[:, 1], c='blue', label='Cities')
    for i, city in enumerate(cities):
        plt.text(city[0], city[1], str(i), fontsize=12)

    # 绘制当前最优路径
    plt.subplot(1, 2, 2)
    path = np.append(best_individual, best_individual[0])
    plt.plot(cities[path, 0], cities[path, 1], 'r-o', label='Best Path')
    plt.title(f'Generation: {generation}, Distance: {1 / best_fitness:.2f}')
    plt.legend()

    plt.tight_layout()
    plt.show()


# 主程序
population = init_population(POP_SIZE, len(cities))
best_individual = None
best_fitness = 0

for generation in range(GEN_MAX):
    # 计算适应度
    fitnesses = [calculate_distance(individual) for individual in population]

    # 更新最优解
    current_best = max(zip(fitnesses, population), key=lambda x: x[0])
    if current_best[0] > best_fitness:
        best_fitness, best_individual = current_best

    # 选择、交叉、变异
    new_population = []
    for _ in range(POP_SIZE // 2):
        parent1 = selection(population, fitnesses)
        parent2 = selection(population, fitnesses)
        child1, child2 = crossover(parent1, parent2), crossover(parent2, parent1)
        mutate(child1)
        mutate(child2)
        new_population.extend([child1, child2])

    population = new_population

    # 可视化每一代的最优解
    if generation % 10 == 0:
        visualize(generation, best_individual, best_fitness)

print(f'最终最优路径：{best_individual}')
print(f'最短距离：{1 / best_fitness:.2f}')